import numpy as np
import torch
import torch.nn.functional as F
from torchvision.models import inception_v3
from torchvision.transforms import ToTensor, Normalize
from torch.utils.data import TensorDataset, DataLoader
from scipy import linalg
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -------- Step 1: Load .npz files --------

# ref_npz = np.load("samples/ref/cifar10.npz")["arr_0"]
# ref_npz = np.load("samples/tanh-SiT-B-2-0040000-size-256-vae-ema-cfg-1.8-seed-0-sde.npz")["arr_0"]
ref_npz = np.load("samples/clean-SiT-B-2-0040000-size-256-vae-ema-cfg-1.8-seed-0-sde.npz")["arr_0"]


gen_npz = np.load("samples/SiT-B-2-0040000-size-256-vae-ema-cfg-1.8-seed-0-sde.npz")["arr_0"]
# gen_npz = np.load("samples/clean-SiT-B-2-0040000-size-256-vae-ema-cfg-1.8-seed-0-sde.npz")["arr_0"]



# Convert to float32 [0, 1]
gen_npz = gen_npz.astype(np.float32) / 255.
ref_npz = ref_npz.astype(np.float32) / 255.

# -------- Step 2: Resize generated images to 32x32 --------
gen_resized = torch.from_numpy(gen_npz).permute(0, 3, 1, 2)  # (N, 3, 256, 256)
gen_resized = F.interpolate(gen_resized, size=(32, 32), mode='bilinear', align_corners=False)

ref_tensor = torch.from_numpy(ref_npz).permute(0, 3, 1, 2)  # (N, 3, 32, 32)

# -------- Step 3: Normalize images as Inception expects --------
def preprocess(images):
    return Normalize([0.5]*3, [0.5]*3)(images)  # scale to [-1, 1]

# -------- Step 4: Extract Inception features --------
def get_inception_features(images, batch_size=64):
    model = inception_v3(pretrained=True, transform_input=False, aux_logits=True)
    model.fc = torch.nn.Identity()  # 提取 pool_3 层（2048-d）特征
    model.eval().to(device)


    dataset = TensorDataset(images)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

    feats = []
    with torch.no_grad():
        for (x,) in tqdm(loader, desc="Extracting features"):
            x = preprocess(x.to(device))
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)  # Resize for Inception
            feat = model(x)
            feats.append(feat.cpu())
    return torch.cat(feats, dim=0)

# Extract features
gen_feats = get_inception_features(gen_resized)
ref_feats = get_inception_features(ref_tensor)

# -------- Step 5: Compute FID --------
def compute_fid(mu1, sigma1, mu2, sigma2):
    """Fréchet Inception Distance between two Gaussians."""
    diff = mu1 - mu2
    covmean, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    fid = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2.0 * covmean)
    return fid

# Convert to numpy
gen_feats_np = gen_feats.numpy()
ref_feats_np = ref_feats.numpy()

mu_gen = np.mean(gen_feats_np, axis=0)
sigma_gen = np.cov(gen_feats_np, rowvar=False)

mu_ref = np.mean(ref_feats_np, axis=0)
sigma_ref = np.cov(ref_feats_np, rowvar=False)

# Compute FID
fid_score = compute_fid(mu_gen, sigma_gen, mu_ref, sigma_ref)
print(f"FID: {fid_score:.4f}")
